import segyio

def find_first_non_zero_time(data, start_index):
    """
    从指定的起始索引开始，查找第一个非零值的时间索引。
    :param data: 数据数组
    :param start_index: 起始索引
    :return: 第一个非零值的时间索引，如果没有非零值，返回None
    """
    for i in range(start_index, len(data)):
        if data[i] != 0:
            return i
    return None

def get_flag(segy_file, start_trace):
    """
        判断sgy文件的数据是向左偏移还是向右偏移
        :param segy_file: SEG-Y文件数据
        :param start_trace: 起始道编号
        :return: 偏移标志，如果是-1，则向左偏移，如果是1，则向右偏移
        """
    flag = -1  # 如果是-1，则向左偏移，如果是1，则向右偏移

    # 从指定道开始处理
    trace_index = start_trace - 1  # segyio的道索引从0开始
    trace_data = segy_file.trace[trace_index]  # 读取当前道的数据
    start_time_index = 0  # 第1道开始时间为0
    non_zero_time_index = find_first_non_zero_time(trace_data, start_time_index)  # 查找第一个非零值的时间索引
    # 左侧第2道数据
    trace_index2 = trace_index + flag  
    trace_data2 = segy_file.trace[trace_index2]  # 读取当前道的数据
    start_time_index2 = non_zero_time_index  # 开始时间为上一道第一个不为0的时间索引
    non_zero_time_index2 = find_first_non_zero_time(trace_data2, start_time_index2)  # 查找第一个非零值的时间索引
    # 判断左偏还是右偏
    if non_zero_time_index == non_zero_time_index2:
        flag = 1
    
    return flag

def get_trace_offset(segy_file, start_trace, trace_num):
    """
    处理SEG-Y文件，从指定道开始，逐道查找第一个非零值的时间。
    :param segy_file: SEG-Y文件数据
    :param start_trace: 起始道编号
    :param trace_num: 走廊通道数
    :return: 道偏移的相关数据，
    [
        {
        "trace_index"，道编号，从1开始
        "non_zero_time"，该道第一个非零点的时间
        },
        ...
    ]
    """
    sample_interval = segy_file.bin[segyio.BinField.Interval] / 1000  # 获取采样间隔（ms），转换为ms
    flag = get_flag(segy_file, start_trace)  # 获取偏移方向
    trace_offset = []  # 存储道偏移的相关数据

    if flag == 1:  # 向右偏移时，开始计算的初始道为最右侧的道
        start_trace = start_trace + trace_num + 1

    # 从指定道开始处理
    trace_index = start_trace + flag  # segyio的道索引从0开始
    last_non_zero_time = 0

    while trace_index >= 0:
        trace_data = segy_file.trace[trace_index]  # 读取当前道的数据
        if trace_index == start_trace - 1:  # 如果是第一道，从0ms开始排查
            start_time_index = 0
        else:  # 否则从上一道记录的时间开始排查
            start_time_index = last_non_zero_time

        # 查找第一个非零值的时间索引
        non_zero_time_index = find_first_non_zero_time(trace_data, start_time_index)
        if non_zero_time_index is not None:  # 有非零值
            non_zero_time = non_zero_time_index * sample_interval  # 通过索引和采样间隔得到事件
            print(f"道 {trace_index + 1}: 第一个非零值的时间为 {non_zero_time} ms")
            last_non_zero_time = non_zero_time_index
            trace_offset.append({
                "trace_index": trace_index + 1,  # 道编号
                "non_zero_time": non_zero_time  # 该道第一个非零点的时间
            })
        else:  # 没有找到非零值
            print(f"道 {trace_index + 1}: 没有找到非零值")
            break

        # 准备处理下一道
        trace_index = trace_index + flag

    return trace_offset

def update_stack_data(file_path, start_trace, trace_num):
    """
    对走廊和滤波走廊数据进行更新，将不属于该走廊部分的数据全部置为0
    :param file_path: SEG-Y文件路径
    :param start_trace: 起始道编号
    :param trace_num: 走廊通道数，原始通道数，不是偏移后的
    :return:
    """
    with segyio.open(file_path, "r+", ignore_geometry=True) as segy_file:
        sample_interval = segy_file.bin[segyio.BinField.Interval] / 1000  # 获取采样间隔（ms），转换为ms
        max_time_index = len(segy_file.trace[start_trace]) - 1  # 计算最大时间索引
        trace_offset = get_trace_offset(segy_file, start_trace, trace_num)  # 获取道偏移的相关数据
        flag = get_flag(segy_file, start_trace)  # 获取数据偏移方向

        # 修改左侧的数据
        for i in range(len(trace_offset)):
            time_index = int(trace_offset[i]["non_zero_time"] / sample_interval)
            trace_index = trace_offset[i]["trace_index"] - 1  # segyio的道索引从0开始
            trace_data = segy_file.trace[trace_index]
            if flag == -1:  # 左偏移
                trace_data[0 : time_index] = 0  # 修改指定范围内的数据
            else:  # 右偏移
                trace_data[time_index: max_time_index + 1] = 0  # 修改指定范围内的数据
            segy_file.trace[trace_index] = trace_data  # 将修改后的数据写回到文件中

        # 修改右侧的数据
        for i in range(len(trace_offset)):
            time_index = int(trace_offset[i]["non_zero_time"] / sample_interval)
            trace_index = trace_offset[i]["trace_index"] - 1 + trace_num # segyio的道索引从0开始
            print(trace_index)
            trace_data = segy_file.trace[trace_index]
            if flag == -1:  # 左偏移
                trace_data[time_index: max_time_index + 1] = 0  # 修改指定范围内的数据
            else:  # 右偏移
                trace_data[0 : time_index] = 0  # 修改指定范围内的数据
            segy_file.trace[trace_index] = trace_data  # 将修改后的数据写回到文件中

def clip_trace(inp_file, out_file, start_index, end_index):
    """
    通过裁剪得到走廊数据，包括start_index和end_index
    :param inp_file: 被截取的sgy文件路径，地面镶嵌数据
    :param out_file: 截取后保存的sgy文件路径，走廊
    :param start_index: 截取道的起始下标，地面镶嵌数据中走廊的起始位置
    :param end_index: 截取道的结束下标，地面镶嵌数据中走廊的结束位置
    """
    with segyio.open(inp_file, 'r', ignore_geometry=True) as infile:
        # 以写入模式打开输出的 SGY 文件
        spec = segyio.spec()
        spec.sorting = infile.sorting
        spec.format = infile.format
        spec.samples = infile.samples
        spec.tracecount = end_index - start_index + 1  # 数据总道数
        with segyio.create(out_file, spec) as outfile:
            # 复制二进制头信息
            outfile.bin = infile.bin
            # 复制道头信息和数据
            for i, trace_index in enumerate(range(start_index - 1, end_index)):
                outfile.header[i] = infile.header[trace_index]
                outfile.trace[i] = infile.trace[trace_index]

def get_st_stkb(file_path, stkbf_path, st_path, stkb_start_trace, trace_num=10):
    """
    将滤波走廊和走廊从sgy文件中分离出来，并保存到指定路径
    :param file_path: sgy文件路径
    :param stkbf_path: 滤波走廊文件保存路径
    :param st_path: 走廊文件保存路径
    :param stkb_start_trace: 滤波走廊的起始道
    :param trace_num: 默认通道数为10
    :return: {
        "trace_start_offset": start_trace2,  # 走廊起始道编号，偏移后的，滤波走廊和走廊相同
        "trace_num_offset": stkb_segy_file.tracecount,  # 总道数，偏移后的，滤波走廊和走廊相同
        "flag": flag,  # 标志，左偏还是右偏
        "trace_offset": trace_offset,  # 偏移相关数据，偏移后的，滤波走廊和走廊相同，左侧第一道为13，右侧第一道为11
        "trace_num": trace_num  # 总道数，偏移前的
    }
    """
    with segyio.open(file_path, "r", ignore_geometry=True) as segy_file:
        flag = get_flag(segy_file, stkb_start_trace)

        # 计算滤波走廊的截取范围
        trace_offset = get_trace_offset(segy_file, stkb_start_trace, trace_num)  # 获取道偏移的相关数据
        if flag == -1:  # 向左偏移
            start_index = trace_offset[len(trace_offset) - 1]["trace_index"]  # 截取道的开始下标，392
            end_index = stkb_start_trace + trace_num  # 截取道的结束下标，414
        else:  # 向右偏移
            start_index = stkb_start_trace + 1  # 截取道的开始下标，405
            end_index = trace_offset[len(trace_offset) - 1]["trace_index"]  # 截取道的结束下标,427
        # 裁剪得到滤波走廊数据
        clip_trace(file_path, stkbf_path, start_index, end_index)
        # 更新滤波走廊数据
        with segyio.open(stkbf_path, "r", ignore_geometry=True) as stkb_segy_file:
            if flag == -1:  # 向左偏移
                start_trace2 = stkb_segy_file.tracecount - trace_num  # 起始道，13
                update_stack_data(stkbf_path, start_trace2, trace_num)
            else:
                start_trace2 = 0  # 起始道，0
                update_stack_data(stkbf_path, start_trace2, trace_num)

        # 计算走廊的截取范围
        start_index2 = start_index + trace_num + 1  # 截取道的开始下标，403，416
        end_index2 = end_index + trace_num + 1  # 截取道的结束下标，425，438
        # 裁剪得到滤波走廊数据
        clip_trace(file_path, st_path, start_index2, end_index2)
        # 更新滤波走廊数据
        with segyio.open(st_path, "r", ignore_geometry=True) as st_segy_file:
            if flag == -1:  # 向左偏移
                start_trace3 = st_segy_file.tracecount - trace_num  # 起始道，13
                update_stack_data(st_path, start_trace3, trace_num)
            else:
                start_trace3 = 0  # 起始道，0
                update_stack_data(st_path, start_trace3, trace_num)

    # 将地面嵌入数据的相关偏移数据，改为裁剪出来的走廊的相关偏移数据
    for i in range(len(trace_offset)):
        trace_offset[i]["trace_index"] = trace_offset[i]["trace_index"] - start_index + 1

    return {
        "trace_start_offset": start_trace2,  # 走廊起始道编号，偏移后的，滤波走廊和走廊相同，左偏为13，右偏为0
        "trace_num_offset": stkb_segy_file.tracecount,  # 总道数，偏移后的，滤波走廊和走廊相同
        "flag": flag,  # 标志，左偏还是右偏
        "trace_offset": trace_offset,  # 偏移相关数据，偏移后的，滤波走廊和走廊相同，左侧第一道为13，右侧第一道为11
        "trace_num": trace_num  # 总道数，偏移前的
    }

if __name__ == "__main__":
    file_path = r"E:\vspData\斜井\psdm_stk_inl994_1394_0914vsp_mu_filter_scale_inq_3D_tm_inlne1194AGC镶嵌1.sgy"
    stkbf_path = r"E:\vspData\斜井\stkb.sgy"  # 滤波走廊
    st_path = r"E:\vspData\斜井\st.sgy"  # 走廊

    # 分离滤波走廊和走廊
    get_st_stkb(file_path, stkbf_path, st_path, 404)

